"""
    生成 zoom circle 图的数据 
"""


import os
import json
import copy
import argparse
import pandas as pd 



def generate_concept_circle(df, tree, tree_concept_all, output, result_tree_name):
    # --------------------- 统计没有子概念的基因簇的数量 -------------------------
    # 统计每个概念出现的数量
    statistic_dict = {}
    for row in df.values.tolist():
        # 提取表格后面六列
        # 6列中第一列是顶级概念，第六列是颗粒度最细的概念
        row_6 = row[-6:]
        if row_6.count("") != 6:
            for index, r in enumerate(row_6):
                if r != "":
                    r_list = r.split("||")
                    for r_i in r_list:
                        concept_info = tree[str(index + 1)][r_i]
                        level = concept_info["level"]
                        if r_i not in statistic_dict:
                            statistic_dict[r_i] = {
                                "num": 0,
                                "level": int(level)
                            }
                        statistic_dict[r_i]["num"] += 1
                                            
    # statistic_dict 转换为 表格
    function_list = []
    for k,v in statistic_dict.items():
        function_list.append([k, v["num"], v["level"]])
    # function_df = pd.DataFrame(function_list, columns=["concept","num","level"])
    # function_df = function_df.sort_values(by="num", ascending=False)
    
    # 提取该物种或染色体概念树中没有子节点的概念
    no_children_list = []
    for f in function_list:
        concept = f[0]
        num = f[1]
        index = f[2]    
        is_exists = ""  # 该标志用于判断有没有子节点概念
        for row in df.values.tolist():  # 遍历表格的每一行，判断 有concept 的行
            row_6 = row[-6:]    # 提取概念树的6列
            if index == 6:      # 如果index是6，那么是没有子节点概念的
                is_exists = ""
            else:
                for r in row_6: # 遍历6列中的单元格
                    r_list = r.split("||")  # 单元格切分 
                    # 如果concept在这行，判断concept在后面几列中有没有子节点
                    if concept in r_list:   
                        next = row_6[index]
                        if row_6[index] != "":
                            # if concept == "Isochorismatase Family":
                            #     print(concept,"\n",row_6,"\n",index,
                            #           "\n",r_list,"\n",next,"\n\n\n")
                            next_parents = []
                            next = next.split("||")
                            for n in next:
                                n_parent = tree[str(index+1)][n]["parent"]
                                next_parents.append(n_parent)
                            if concept in next_parents:
                                is_exists = True
        if is_exists == True:
            # print(concept, "\t", num, "\t", index)
            pass 
        else:
            # print("\t\t".join([str(x) for x in f]))
            no_children_list.append(f)
            
    # no_children_df = pd.DataFrame(no_children_list, columns=["funciont","num","level"])
    # no_children_df.to_excel("test.xlsx", index=None)
    
    # 没有子节点的概念
    no_children_dict = {}
    for row in no_children_list:
        no_children_dict[row[0]] = {
            "num": row[1],
            "level": row[2]
        }
   
    # ------------------------ 生成概念树每层的数据 ------------------------------
    # 数树一个共多少层
    depth = 0
    species_concepts = {}
    for k,v in tree_concept_all.items():
        level_list = v["level_list"]
        # print(v,"\t", level_list)
        if isinstance(level_list, list):
            length = len(level_list)
            if length > depth:
                depth = length
        concepts = v["concepts"]
        for c in concepts:
            if c not in species_concepts:
                species_concepts[c] = ""
    # print(depth)
    # print(len(species_concepts))
    # 生成每层的数据
    layer_dict = {
        1: [],
        2: [],
        3: [],
        4: [],
        5: [],
        6: []
    }
    for c in species_concepts:
        c_info = {}
        for k,v in tree.items():
            if c in v:
                c_info = v[c]
        if c_info != {}:
            c_level = int(c_info["level"])
            if c in no_children_dict:
                layer_dict[c_level].append({
                    "name": c,
                    "value": no_children_dict[c]["num"]
                })
            else:
                layer_dict[c_level].append({
                    "name": c,
                    "children": []
                })

    # ----------------- 把层与层之间关联起来成为 zoom circle图的数据 --------------
    tmp_tree_json = {}
    layer_range = list(range(2,depth+1))[::-1]
    for layer_index, i in enumerate(layer_range):
        if len(layer_range) > 1: 
            if i == layer_range[0]: # 最细一层单独处理
                tmp_tree_json[i] = []
                layer_start = layer_dict[i]
                layer_tmp_dict = {}
                for l in layer_start:
                    name = l["name"]
                    parent = tree[str(i)][name]["parent"]
                    if parent not in layer_tmp_dict:
                        layer_tmp_dict[parent] = []
                    layer_tmp_dict[parent].append(l)
                for l_k,l_v in layer_tmp_dict.items():
                    tmp_tree_json[i].append({
                        "name": l_k,
                        "children": l_v
                    })
            else:
                tmp_tree_json[i] = []
                layer = copy.deepcopy(layer_dict[i])  # 当前处理的那一层
                layer_more_1 = tmp_tree_json[i+1]        # 获取当前处理的子一层数据
                layer_tmp_dict = {}
                for l in layer:
                    if "children" in l:
                        name = l["name"] 
                        l_v = []
                        for l_m_1 in layer_more_1:
                            l_m_name = l_m_1["name"]
                            if name == l_m_name:
                                l_v = l_m_1["children"]
                        parent = tree[str(i)][name]["parent"]
                        if parent not in layer_tmp_dict:
                            layer_tmp_dict[parent] = []
                        layer_tmp_dict[parent].append({
                            "name": name,
                            "children": l_v
                        })
                    else:
                        name = l["name"]
                        parent = tree[str(i)][name]["parent"]
                        if parent not in layer_tmp_dict:
                            layer_tmp_dict[parent] = []
                        layer_tmp_dict[parent].append(l) 
                for l_k,l_v in layer_tmp_dict.items():
                    tmp_tree_json[i].append({
                        "name": l_k,
                        "children": l_v
                    })
        elif len(layer_range) == 1:
            if i == layer_range[0]: # 最细一层单独处理
                tmp_tree_json[i] = []
                layer_start = layer_dict[i]
                layer_tmp_dict = {}
                for l in layer_start:
                    name = l["name"]
                    parent = tree[str(i)][name]["parent"]
                    if parent not in layer_tmp_dict:
                        layer_tmp_dict[parent] = []
                    layer_tmp_dict[parent].append(l)
                for l_k,l_v in layer_tmp_dict.items():
                    tmp_tree_json[i].append({
                        "name": l_k,
                        "children": l_v
                    })
    if depth == 0:
        tmp_tree_json[2] = {
            "name": tree_concept_all,
            "children": []
        }
    
    if 2 in tmp_tree_json:
        reuslt_dict = {
            "name": result_tree_name,
            "children": tmp_tree_json[2]
        }
        with open(output, mode="w", encoding="utf-8") as file:
            json.dump(reuslt_dict, file, indent=4)
        
        
        
def main(annotation_table, tree, concept_parent_all, output_dir):
    # 读入概念树注释结果的表格
    all_df = pd.read_excel(annotation_table, keep_default_na=False)
    # 需要使用染色体名加上簇名一起对表格进行去重，因为这里只需要每一个簇的注释结果
    # df = df.drop_duplicates(["chromosome/scaffold/contif","cluster_name"])
    title = all_df.columns.tolist()
    # print(title)
    
    # 读入树
    with open(tree, encoding="utf-8") as file:
        tree = json.load(file)
    
    # 读入每个物种标签对上概念树中的所有标签的数据
    with open(concept_parent_all, encoding="utf-8") as file:
        tree_concept_all = json.load(file)
    
    # 生成每个染色体对应的概念、表格和json数据
    groups_chrom = all_df.groupby("chromosome/scaffold/contif")
    chrom_dict = {}
    for key, sub_df in groups_chrom:
        if key not in chrom_dict:
            chrom_dict[key] = {
                "df": sub_df,
                "chrom_json": {}
            }
        chrom_json = {}
        for row in sub_df.values.tolist():
            row_7 = row[-7]
            summary_feature = eval(
                "['" + row_7.strip('"]').strip("'").strip("['").strip('"').replace(
                "', '", "---").replace('"', "'").replace("'", "").replace("---", "', '")
                + "']"
            )
            for f in summary_feature:
                if f in tree_concept_all:
                    chrom_json[f] = tree_concept_all[f]
        chrom_dict[key]["chrom_json"] = chrom_json
        
    file_name = os.path.basename(annotation_table).split(".")[0]
    output_dir = os.path.join(output_dir, file_name)
    os.makedirs(output_dir, exist_ok=True)
    
    tasks = []
    for k,v in chrom_dict.items():
        output = os.path.join(output_dir, k + ".json")
        df = v["df"]
        chrom_json = v["chrom_json"]
        tasks.append([df, tree, chrom_json, output, k])
        
    tasks.append([all_df, tree, tree_concept_all, os.path.join(output_dir, "all.json"), file_name])
        
    for task in tasks[0:]:
        generate_concept_circle(task[0], task[1], task[2], task[3], task[-1])


def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("-annotation_table", "--annotation_table")
    parser.add_argument("-tree", "--tree")
    parser.add_argument("-concept_parent_all", "--concept_parent_all")
    parser.add_argument("-output_dir", "--output_dir")
    params = parser.parse_args()
    
    return params 


if __name__ == "__main__":
    params = parse_args()
    annotation_table = params.annotation_table
    tree = params.tree
    concept_parent_all = params.concept_parent_all
    output_dir = params.output_dir
    
    # annotation_table = "09.cluster_tree_function_annotation.xlsx"
    # tree = "16.menu_tree_to_dict_embedding.json"
    # concept_parent_all = "07.label_to_tree_concept_all.json"
    # output_dir = "10.concept_circle"
    
    main(annotation_table, tree, concept_parent_all, output_dir)
    